package cn._51doit.flink.day07;

import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * 将窗口内的数据进行增量聚合并且与历史数据进行聚合
 * 不使用reduce方法，reduce方法不够灵活，因为输入的类型和返回的类型必须一致
 *
 * aggregate方法更灵活
 */
public class AggregateWindowFunction {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        //1000,spark,3
        //3000,hadoop,5
        DataStreamSource<String> lines = env.socketTextStream("localhost", 8888);

        //提取数据中的eventTime
        SingleOutputStreamOperator<String> dataStreamWithWaterMark = lines.assignTimestampsAndWatermarks(
                WatermarkStrategy.<String>forBoundedOutOfOrderness(Duration.ofSeconds(0)).withTimestampAssigner(new SerializableTimestampAssigner<String>() {
                    @Override
                    public long extractTimestamp(String ele, long l) {
                        String[] fields = ele.split(",");
                        return Long.parseLong(fields[0]);
                    }
                }));

        SingleOutputStreamOperator<Tuple2<String, Integer>> wordAndCount = dataStreamWithWaterMark.map(new MapFunction<String, Tuple2<String, Integer>>() {
            @Override
            public Tuple2<String, Integer> map(String line) throws Exception {
                String[] fields = line.split(",");
                String word = fields[1];
                int count = Integer.parseInt(fields[2]);
                return Tuple2.of(word, count);
            }
        });

        //先KeyBy
        KeyedStream<Tuple2<String, Integer>, String> keyed = wordAndCount.keyBy(tp -> tp.f0);

        //按照EventTime滚动的窗口，窗口长度为10秒
        WindowedStream<Tuple2<String, Integer>, String, TimeWindow> window = keyed
                .window(TumblingEventTimeWindows.of(Time.seconds(10)));

        //sum只会对当前窗口内的数据进行聚合，不会聚合历史数据
        //SingleOutputStreamOperator<Tuple2<String, Integer>> res = window.sum(1);
        SingleOutputStreamOperator<Tuple2<String, Integer>> res = window.aggregate(new MyAggregateFunction(), new MyProcessWindowFunction());

        res.print();

        env.execute();

    }

    private static class MyAggregateFunction implements AggregateFunction<Tuple2<String, Integer>, Integer, Integer> {

        //创建每个窗口每个key的初始值
        @Override
        public Integer createAccumulator() {
            return 0;
        }

        //每输入一条数据调用一次add方法
        @Override
        public Integer add(Tuple2<String, Integer> value, Integer accumulator) {
            return accumulator + value.f1;
        }

        //返回当前窗口，当前key聚合的结果
        @Override
        public Integer getResult(Integer accumulator) {
            return accumulator;
        }

        //该方法只有session window可能会调用，其他类型的窗口不调用
        @Override
        public Integer merge(Integer a, Integer b) {
            return null;
        }
    }


    private static class MyProcessWindowFunction extends ProcessWindowFunction<Integer, Tuple2<String, Integer>, String, TimeWindow> {

        private transient ValueState<Integer> countState;

        @Override
        public void open(Configuration parameters) throws Exception {
            //初始化或恢复状态
            //定义状态描述器（描述状态的名称和类型）
            ValueStateDescriptor<Integer> stateDescriptor = new ValueStateDescriptor<>("count-state", Integer.class);
            //通过运行时上下文初始化或恢复状态
            countState = getRuntimeContext().getState(stateDescriptor);
        }

        //窗口触发后会调用该方法
        //会将增量聚合后的结果传入到该方法中，窗口触发后每个key都会调用一次该方法
        @Override
        public void process(String key, Context context, Iterable<Integer> elements, Collector<Tuple2<String, Integer>> out) throws Exception {
            //当前窗口增量聚合的结果
            Integer currentCount = elements.iterator().next();
            //跟历史数据进行累加
            Integer historyCount = countState.value();
            if (historyCount == null) {
                historyCount = 0;
            }
            historyCount += currentCount;
            //更新状态
            countState.update(historyCount);
            //输出
            out.collect(Tuple2.of(key, historyCount));
        }
    }

}
